package com.example.studyllm.component.service;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.example.studyllm.bo.ModelBO;
import com.example.studyllm.component.model.OllamaChatModelBuilderHandler;
import com.example.studyllm.component.model.OpenAiChatModelBuilderHandler;
import com.example.studyllm.component.model.QwenChatModelBuilderHandler;
import com.example.studyllm.constants.ModelSupplierConstant;
import com.example.studyllm.entity.rag.RagModel;
import com.example.studyllm.service.rag.RagModelService;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author wk
 */
@Slf4j
@Component
@AllArgsConstructor
public class ModelInitComponent implements ApplicationContextAware {

    private final RagModelService ragModelService ;
    public void init() {
        List<RagModel> ragModels = ragModelService.list();
        if(CollectionUtil.isEmpty(ragModels)){
            return ;
        }
        ragModels.forEach(ragModel -> {
            ModelBO modelBO = ModelBO.builder()
                    .id(ragModel.getCode())
                    .name(ragModel.getName())
                    .code(ragModel.getCode())
                    .top(ragModel.getTop())
                    .apiKey(ragModel.getApiKey())
                    .type(ragModel.getType())
                    .baseUrl(ragModel.getUrl())
                    .model(ragModel.getModelName())
                    .dimension(ragModel.getDimensions())
                    .responseLength(ragModel.getResponseLength())
                    .secretKey(ragModel.getSecretKey())
                    .supplier(ragModel.getSupplier())
                    .temperature(ragModel.getTemperature())
                    .build();
            if(ragModel.getType().equals("MODEL")){
                if(ragModel.getSupplier().equals(ModelSupplierConstant.OLLAMA)){
                    SpringUtil.registerBean(ragModel.getCode(), new OllamaChatModelBuilderHandler().buildModel(modelBO));
                }else if(ragModel.getSupplier().equals(ModelSupplierConstant.OPENAI)){
                    SpringUtil.registerBean(ragModel.getCode(), new OpenAiChatModelBuilderHandler().buildModel(modelBO));
                }else if(ragModel.getSupplier().equals(ModelSupplierConstant.QWEN)){
                    SpringUtil.registerBean(ragModel.getCode(), new QwenChatModelBuilderHandler().buildModel(modelBO));
                }
            }else{
                if(ragModel.getSupplier().equals(ModelSupplierConstant.OLLAMA)){
                    SpringUtil.registerBean(ragModel.getCode(), new OllamaChatModelBuilderHandler().buildEmbeddingModel(modelBO));
                }else if(ragModel.getSupplier().equals(ModelSupplierConstant.OPENAI)){
                    SpringUtil.registerBean(ragModel.getCode(), new OpenAiChatModelBuilderHandler().buildEmbeddingModel(modelBO));
                }else if(ragModel.getSupplier().equals(ModelSupplierConstant.QWEN)){
                    SpringUtil.registerBean(ragModel.getCode(), new QwenChatModelBuilderHandler().buildEmbeddingModel(modelBO));
                }
            }

        });
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        init();
    }
}
